package cn.lisens.demo.service;

import org.springframework.http.HttpStatus;
import org.springframework.util.ObjectUtils;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.ExchangeFunction;
import reactor.core.publisher.Mono;

public class AuthorizationFilterFunction implements ExchangeFilterFunction {
    private String token;
    private AuthorizationClient authorizationClient;
    public AuthorizationFilterFunction(final AuthorizationClient authorizationClient) {
        this.authorizationClient = authorizationClient;
    }

    @Override
    public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
        if (ObjectUtils.isEmpty(token)) {
            return this.throughGetToken(request, next);
        }
        return withAuthorizationHeader(request, next);
    }

    private Mono<ClientResponse> throughGetToken(ClientRequest request, ExchangeFunction next) {
        return this.authorizationClient.getToken()
                .flatMap(response -> {
                    this.token = response;
                    return withAuthorizationHeader(request, next);
                });
    }

    private Mono<ClientResponse> withAuthorizationHeader(ClientRequest request, ExchangeFunction next) {
        final ClientRequest withAuthorizationHeaderRequest = ClientRequest.from(request).header("Authorization", "Bearer " + this.token)
                .build();
        return next.exchange(withAuthorizationHeaderRequest)
                .flatMap(clientResponse -> {
                    if (clientResponse.statusCode() == HttpStatus.UNAUTHORIZED) {
                         return throughGetToken(request, next);
                    }
                    return Mono.just(clientResponse);
                });
    }
}
